import os.path as osp
from shutil import copyfile
import time
import requests
from typing import Dict, List, Optional

from pydantic import BaseModel
import yaml
from jinja2 import Template

from tjob.configs.cloud import CLOUD_DEFAULT_CONF_PATH, CLOUD_DEFAULT_JOB_TEMPLATE_PATH, CLOUD_DEFAULT_BATCH_JOB_TEMPLATE_PATH, CLOUD_DEFALT_ACCESS_PATH, CLOUD_DEFAULT_JOB_PARAMS_PATH
from tjob.initialize.setup_envs import _CACHE
from tjob.utils.io_utils import save_yaml, load_yaml


_CLOUD_ACCESS = osp.join(_CACHE, "cloud.access.yaml")
_CLOUD_JOB_TEMPLATE = osp.join(_CACHE, "cloud.job.template.yaml")
_CLOUD_BATCH_JOB_TEMPLATE = osp.join(_CACHE, "cloud.batch.job.template.yaml")
_CLOUD_JOB_PARAMS = osp.join(_CACHE, "cloud.job.params.yaml")
_CLOUD_DEFAULT_CONF = osp.join(_CACHE, "cloud.config.yaml")


class CloudAccess(BaseModel):
    url: str
    username: str
    password: str
    token: str
    webURL: str
    refreshToken: str

class Tolerations(BaseModel):
    key: str
    operator: str
    value: str
    effect: str


class GPUConfig(BaseModel):
    cpu: int
    memory: int
    nodeSelector: Dict[str,str]
    tolerations: List[Tolerations]

class GPUMapping(BaseModel):
    gpu: Dict[str, GPUConfig]

class Config(BaseModel):
    gpu: Optional[GPUMapping]
    randomWhenConflict: bool
    cloudCluster: str
    cloudWorkspace: str

class VolumeMount(BaseModel):
    name: str
    mountPath: str
    
class VolumeOption(BaseModel):
    key:str
    value: str

class Volume(BaseModel):
    name: str
    volumeType: str
    volumeOptions: List[VolumeOption]

class Jupyter(BaseModel):
    enabled: bool = False
    workingDir: str = "/"
    ingressClass: str = "traefik"
    port: int = 8888
    host: str # = "lab-jobsscz.srv.deeproute.cn"
    autoInstall: bool = False

class Batch(BaseModel):
    minAvailable: int = 1
    minSuccess: int = 0
    restartPolicy: str = "Never"

class Framework(BaseModel):
    batch: Batch = None

PRIORITY_CLASS = {
    1: "pai-low",
    2: "pai-middle",
    3: "pai-high",
    4: "pai-urgent",
}

class CloudJobParams(BaseModel):
    cloudCluster:str
    name: str
    image: str
    command: str
    count: int
    gpuType: str
    gpuPerWorker: int = 0
    cpuPerWorker: int
    memPerWorker: int
    priority: int = 1
    shmSize:  int
    volumes: List[Volume] 
    volumeMounts: List[VolumeMount] 
    jupyter: Jupyter = None
    framework: Framework = None
    isDevJob: str

    @property
    def tolerations(self):
        pass

    def to_cloud_job(self) ->dict:
        if self.framework and self.framework.batch:
            job_template = Template(get_cloud_batch_job_template())
        else:
            job_template = Template(get_cloud_job_template())
        default_conf = get_cloud_default_conf()
        
        nodeSelector = default_conf["gpu"][self.gpuType]["nodeSelector"]
        tolerations = default_conf["gpu"][self.gpuType]["tolerations"]
        priority_class = PRIORITY_CLASS.get(self.priority, "pai-low")
        access = get_cloud_access()
        self.gpuPerWorker = 0 if self.gpuType == "none" else self.gpuPerWorker
        job_str = job_template.render(
            user = access["username"],
            namespace = get_default_namespace(access),
            name = self.name,
            count = self.count,
            image = self.image,
            command = self.command,
            gpuPerWorker = self.gpuPerWorker,
            cpuPerWorker = self.cpuPerWorker,
            memPerWorker = self.memPerWorker,
            workerCount = self.count - 1,
            shmSize = self.shmSize,
            volumes = self.volumes,
            volumeMounts = self.volumeMounts,
            nodeSelector=nodeSelector,
            tolerations=tolerations,
            priorityClass=priority_class,
            jupyter=self.jupyter,
            framework=self.framework,
            isDevJob=self.isDevJob,
        )

        return yaml.safe_load(job_str)

def dump_cloud_access(userame, password, token):
    conf = get_cloud_access()
    conf["username"] = userame
    conf["password"] = password
    conf["token"] = token
    save_yaml(_CLOUD_ACCESS, conf)


def get_cloud_job_template() -> str:
    if not osp.exists(_CLOUD_JOB_TEMPLATE):
        copyfile(CLOUD_DEFAULT_JOB_TEMPLATE_PATH, _CLOUD_JOB_TEMPLATE)
    with open(_CLOUD_JOB_TEMPLATE, "r") as f:
        return f.read()

def get_cloud_batch_job_template() -> str:
    if not osp.exists(_CLOUD_BATCH_JOB_TEMPLATE):
        copyfile(CLOUD_DEFAULT_BATCH_JOB_TEMPLATE_PATH, _CLOUD_BATCH_JOB_TEMPLATE)
    with open(_CLOUD_BATCH_JOB_TEMPLATE, "r") as f:
        return f.read()

def get_cloud_access() -> dict:
    if not osp.exists(_CLOUD_ACCESS):
        copyfile(CLOUD_DEFALT_ACCESS_PATH, _CLOUD_ACCESS)
    return load_yaml(_CLOUD_ACCESS)

def set_cloud_login(username, password) -> str:
    conf = get_cloud_access()
    url = conf["url"] + "/oauth/keycloak-token"
    data = {"username": username, "password": password, "grant_type": "password"}
    headers = {
    "Content-Type": "application/x-www-form-urlencoded"
    }
    res = requests.post(url, data=data, headers=headers)
    if res.status_code != 200:
        raise ValueError(f"Fetch cloud token failed: {res.text}")
    token = res.json()["access_token"]
    conf["username"] = username
    conf["password"] = password
    conf["token"] = token
    dump_cloud_access(username, password, token)

    init_except=None
    for _ in range(3):
        try:
            _init_cloud_user_namespace(conf)
            _init_cloud_rasmap_pvc(conf)
            break
        except Exception as e:
            init_except = e
            time.sleep(1)
    else:
        raise init_except
    return token



def get_default_namespace(access):
    return access["username"] +"-training"

def set_cloud_default_params(params):
    default_params = get_cloud_job_default_params()
    if not params:
        return
    default_params.update(params)
    save_yaml(_CLOUD_JOB_PARAMS, default_params)

def get_cloud_job_default_params() -> dict:
    if not osp.exists(_CLOUD_JOB_PARAMS):
        copyfile(CLOUD_DEFAULT_JOB_PARAMS_PATH, _CLOUD_JOB_PARAMS)
    return load_yaml(_CLOUD_JOB_PARAMS)

def get_cloud_default_conf():
    if not osp.exists(_CLOUD_DEFAULT_CONF):
        copyfile(CLOUD_DEFAULT_CONF_PATH, _CLOUD_DEFAULT_CONF)
    return load_yaml(_CLOUD_DEFAULT_CONF)


def _init_cloud_user_namespace(conf):
    config = get_cloud_default_conf()
    url = conf["url"] + f"/kapis/clusters/{config['cloudCluster']}/tenant.kubesphere.io/v1alpha2/workspaces/{config['cloudWorkspace']}/namespaces"
    token = conf["token"]
    if not token:
        token = set_cloud_login(conf["username"], conf["password"])
    import requests
    data = {
        "apiVersion":"v1",
        "kind":"Namespace",
        "metadata":{
            "name":get_default_namespace(conf),
            "labels":{
                "kubesphere.io/workspace":config['cloudWorkspace'],
            },
            "annotations":{
                "kubesphere.io/creator":conf['username'],
            }
        },
        "cluster": config['cloudCluster'],
    }
    res = requests.post(url, json=data,headers={"Authorization": f"Bearer {token}"})
    if res.ok or res.status_code == 400:
        return
    raise ValueError(f"Init cloud user namespace failed: {res.status_code}, {res.text}")


def _init_cloud_rasmap_pvc(access):
    config = get_cloud_default_conf()
    token = access["token"] if access["token"] else set_cloud_login(access["username"], access["password"])
    pvc_str = f"""
apiVersion: v1
kind: PersistentVolumeClaim
metadata:
  annotations:
    kubesphere.io/creator: leili
  finalizers:
    - kubernetes.io/pvc-protection
  labels:
    ks_workspace: {config["cloudWorkspace"]}
  name: rasmap
  namespace: {get_default_namespace(access)}
spec:
  accessModes:
    - ReadOnlyMany
  resources:
    requests:
      storage: 10Gi
  storageClassName: nfs-yanrong-rasmap
"""
    pvc = yaml.safe_load(pvc_str)
    pvc_url = access["url"] + f"/api/clusters/{config['cloudCluster']}/v1/namespaces/{get_default_namespace(access)}/persistentvolumeclaims"
    res = requests.post(pvc_url, json=pvc,headers={"Authorization": f"Bearer {token}"})
    if res.ok or res.status_code == 409: return
    raise ValueError(f"Init cloud user rasmap pv failed: {res.status_code}, {res.text}")

def cloud_job_detail_url(cj:CloudJobParams):
    config = get_cloud_default_conf()
    access = get_cloud_access()
    return access["url"] +f"/{config['cloudWorkspace']}/clusters/{cj.cloudCluster}/projects/{get_default_namespace(config)}/applications/training-tasks/{cj.name}"